# %%
import os
from pathlib import Path

print('workdir:', os.getcwd())
# os.chdir()
PATH = Path(r'C:\files\git_repository\pytorch-learning\2pytorch基础入门实战\11基于全连接神经网络的数字识别')
print('PATH:', PATH)

# %%
import os
import torch
from torch import nn
import torch.nn.functional as F
from torchvision import transforms
from torchvision.datasets import MNIST
from torch.utils.data import DataLoader, random_split
import pytorch_lightning as pl
import torchvision
from torchvision import transforms


# %%
# 定义模型
class LitMNIST(pl.LightningModule):
    def __init__(self):
        super(LitMNIST, self).__init__()
        self.linear1 = nn.Linear(in_features=28 * 28 * 1, out_features=64)
        self.linear2 = nn.Linear(in_features=64, out_features=10)
        # 最后一层不加softmax，后面直接用cross_entropy

    def forward(self, x):
        x = self.linear1(x)
        x = self.relu(x)
        x = self.linear2(x)
        return x

    def training_step(self, batch, batch_idx):
        x, y = batch
        y_pred = self.forward(x)
        loss = F.cross_entropy(y_pred, y)
        self.log('train_loss', loss)
        return loss

    def test_step(self, batch, batch_idx):
        x, y = batch
        y_pred = self.forward(x)
        loss = F.cross_entropy(y_pred, y)
        self.log('train_loss', loss)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
        return optimizer


# %%
# 数据
trans = transforms.Compose([transforms.ToTensor()])
train_dataset = torchvision.datasets.MNIST(root=str(PATH / 'data'), train=True, transform=trans, download=False)
test_dataset = torchvision.datasets.MNIST(root=str(PATH / 'data'), train=False, transform=trans, download=False)
batch_size = 100
# 根据数据集定义数据加载器
train = torch.utils.data.DataLoader(dataset=train_dataset,
                                    batch_size=batch_size,
                                    shuffle=True)

test = torch.utils.data.DataLoader(dataset=test_dataset,
                                   batch_size=batch_size,
                                   shuffle=False)
train, test
# %%
train = iter(train)
example_x, example_y = next(train)
print(example_x.shape)
print(example_y.shape)
# %%
net = LitMNIST()
trainer = pl.Trainer(max_epochs=3, progress_bar_refresh_rate=20)
trainer.fit(net, train)
